#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
机械臂控制功能函数
Robot Arm Control Functions

为LLM提供机械臂控制的Function Calling接口
基于现有的RobotArm实现
"""

import logging
import time
from typing import Dict, Any, Optional, Union, Tuple

from ..servo_arm_control.arm_move import RobotArm
from .function_registry import FunctionRegistry

logger = logging.getLogger(__name__)

# 机械臂限位设置 (单位: 毫米)
ARM_LIMIT_X_MIN = 50
ARM_LIMIT_X_MAX = 180
ARM_LIMIT_Y_MIN = -80
ARM_LIMIT_Y_MAX = 80
ARM_LIMIT_Z_MIN = -50
ARM_LIMIT_Z_MAX = 140

# 全局机械臂控制器实例
_arm_controller: Optional[RobotArm] = None

def get_arm_controller() -> RobotArm:
    """获取机械臂控制器实例"""
    global _arm_controller
    if _arm_controller is None:
        raise RuntimeError("机械臂控制器未初始化，请先调用initialize_arm_controller函数")
    return _arm_controller

def validate_arm_position(x: float, y: float, z: float) -> tuple[bool, str]:
    """
    验证机械臂位置是否在限位范围内
    
    Args:
        x: X坐标 (毫米)
        y: Y坐标 (毫米)
        z: Z坐标 (毫米)
        
    Returns:
        tuple: (是否有效, 错误信息)
    """
    if not (ARM_LIMIT_X_MIN <= x <= ARM_LIMIT_X_MAX):
        return False, f"机械臂X坐标超出限位范围 {x}mm (有效范围: {ARM_LIMIT_X_MIN}-{ARM_LIMIT_X_MAX}mm)"
    if not (ARM_LIMIT_Y_MIN <= y <= ARM_LIMIT_Y_MAX):
        return False, f"机械臂Y坐标超出限位范围 {y}mm (有效范围: {ARM_LIMIT_Y_MIN}-{ARM_LIMIT_Y_MAX}mm)"
    if not (ARM_LIMIT_Z_MIN <= z <= ARM_LIMIT_Z_MAX):
        return False, f"机械臂Z坐标超出限位范围 {z}mm (有效范围: {ARM_LIMIT_Z_MIN}-{ARM_LIMIT_Z_MAX}mm)"
    return True, ""

def initialize_arm_controller(port: str = '/dev/ttyS1', baudrate: int = 115200, auto_init: bool = False) -> bool:
    """
    初始化机械臂控制器
    
    Args:
        port: 串口设备路径
        baudrate: 波特率
        auto_init: 是否自动执行初始化动作
        
    Returns:
        bool: 是否成功初始化
    """
    global _arm_controller
    
    try:
        # 创建机械臂控制器实例，默认不执行自动初始化
        _arm_controller = RobotArm(auto_init=auto_init)
        logger.info(f"机械臂控制器初始化成功")
        return True
    except Exception as e:
        logger.error(f"机械臂控制器初始化失败: {e}")
        _arm_controller = None
        return False

@FunctionRegistry.register_function('arm_catch_object', '控制机械臂抓取物体')
def arm_catch_object_func(sth_position_x: float, sth_position_y: float, sth_position_z: float = -50.0) -> str:
    """
    控制机械臂抓取物体
    
    Args:
        sth_position_x: 物体在图像中的x坐标
        sth_position_y: 物体在图像中的y坐标
        sth_position_z: 抓取高度，默认-50.0
        
    Returns:
        str: 执行结果描述
    """
    try:
        # 参数验证
        if not isinstance(sth_position_x, (int, float)) or not isinstance(sth_position_y, (int, float)) or not isinstance(sth_position_z, (int, float)):
            return "错误：坐标参数必须是数字"
        
        # 限位检查
        is_valid, error_msg = validate_arm_position(sth_position_x, sth_position_y, sth_position_z)
        if not is_valid:
            return f"限位错误：{error_msg}。请调整目标位置到安全范围内。"
        
        controller = get_arm_controller()
        
        # 执行抓取动作
        controller.arm_catch_sth(sth_position_x, sth_position_y, sth_position_z)
        
        return f"机械臂抓取动作执行成功：位置({sth_position_x}, {sth_position_y}, {sth_position_z})"
            
    except RuntimeError as e:
        return f"错误：{str(e)}"
    except Exception as e:
        logger.error(f"机械臂抓取异常: {e}")
        return f"机械臂抓取过程中出现错误：{str(e)}"

@FunctionRegistry.register_function('arm_put_object', '控制机械臂放置物体')
def arm_put_object_func(sth_position_x: float, sth_position_y: float, sth_position_z: float = -50.0) -> str:
    """
    控制机械臂放置物体
    
    Args:
        sth_position_x: 放置位置的x坐标
        sth_position_y: 放置位置的y坐标
        sth_position_z: 放置高度，默认-50.0
        
    Returns:
        str: 执行结果描述
    """
    try:
        # 参数验证
        if not isinstance(sth_position_x, (int, float)) or not isinstance(sth_position_y, (int, float)) or not isinstance(sth_position_z, (int, float)):
            return "错误：坐标参数必须是数字"
        
        # 限位检查
        is_valid, error_msg = validate_arm_position(sth_position_x, sth_position_y, sth_position_z)
        if not is_valid:
            return f"限位错误：{error_msg}。请调整目标位置到安全范围内。"
        
        controller = get_arm_controller()
        
        # 执行放置动作
        controller.arm_normal_put_sth(sth_position_x, sth_position_y, sth_position_z)
        
        return f"机械臂放置动作执行成功：位置({sth_position_x}, {sth_position_y}, {sth_position_z})"
            
    except RuntimeError as e:
        return f"错误：{str(e)}"
    except Exception as e:
        logger.error(f"机械臂放置异常: {e}")
        return f"机械臂放置过程中出现错误：{str(e)}"

@FunctionRegistry.register_function('arm_stack_put_object', '控制机械臂堆叠放置物体')
def arm_stack_put_object_func(sth_position_x: float, sth_position_y: float, sth_position_z: float = 30.0) -> str:
    """
    控制机械臂堆叠放置物体
    
    Args:
        sth_position_x: 放置位置的x坐标
        sth_position_y: 放置位置的y坐标
        sth_position_z: 堆叠放置高度，默认30.0
        
    Returns:
        str: 执行结果描述
    """
    try:
        # 参数验证
        if not isinstance(sth_position_x, (int, float)) or not isinstance(sth_position_y, (int, float)) or not isinstance(sth_position_z, (int, float)):
            return "错误：坐标参数必须是数字"
        
        # 限位检查
        is_valid, error_msg = validate_arm_position(sth_position_x, sth_position_y, sth_position_z)
        if not is_valid:
            return f"限位错误：{error_msg}。请调整目标位置到安全范围内。"
        
        controller = get_arm_controller()
        
        # 执行堆叠放置动作
        controller.arm_stack_put_sth(sth_position_x, sth_position_y, sth_position_z)
        
        return f"机械臂堆叠放置动作执行成功：位置({sth_position_x}, {sth_position_y}, {sth_position_z})"
            
    except RuntimeError as e:
        return f"错误：{str(e)}"
    except Exception as e:
        logger.error(f"机械臂堆叠放置异常: {e}")
        return f"机械臂堆叠放置过程中出现错误：{str(e)}"

@FunctionRegistry.register_function('arm_init_action', '执行机械臂初始化动作')
def arm_init_action_func() -> str:
    """
    执行机械臂初始化动作
    
    Returns:
        str: 执行结果描述
    """
    try:
        controller = get_arm_controller()
        
        # 执行初始化动作
        controller.arm_init_action()
        
        return "机械臂初始化动作执行成功"
            
    except RuntimeError as e:
        return f"错误：{str(e)}"
    except Exception as e:
        logger.error(f"机械臂初始化异常: {e}")
        return f"机械臂初始化过程中出现错误：{str(e)}"

@FunctionRegistry.register_function('arm_move_direction', '控制机械臂向指定方向移动')
def arm_move_direction_func(direction: str, distance: float) -> str:
    """
    控制机械臂向指定方向移动
    
    Args:
        direction: 移动方向 (forward/back/left/right/up/down)
        distance: 移动距离（厘米）
        
    Returns:
        str: 执行结果描述
    """
    try:
        controller = get_arm_controller()
        
        # 参数验证
        if not isinstance(distance, (int, float)) or distance <= 0:
            return "错误：距离必须是正数"
        
        direction_map = {
            "forward": controller.arm_move_forward,
            "back": controller.arm_move_back,
            "left": controller.arm_move_left,
            "right": controller.arm_move_right,
            "up": controller.arm_move_up,
            "down": controller.arm_move_down
        }
        
        if direction not in direction_map:
            return f"错误：无效的方向 '{direction}'，支持的方向: {list(direction_map.keys())}"
        
        # 获取当前位置并计算目标位置
        current_x, current_y, current_z = controller.get_current_position()
        
        # 将距离从厘米转换为毫米
        distance_mm = distance * 10
        
        # 计算目标位置
        if direction == "forward":
            target_x, target_y, target_z = current_x + distance_mm, current_y, current_z
        elif direction == "back":
            target_x, target_y, target_z = current_x - distance_mm, current_y, current_z
        elif direction == "left":
            target_x, target_y, target_z = current_x, current_y - distance_mm, current_z
        elif direction == "right":
            target_x, target_y, target_z = current_x, current_y + distance_mm, current_z
        elif direction == "up":
            target_x, target_y, target_z = current_x, current_y, current_z + distance_mm
        elif direction == "down":
            target_x, target_y, target_z = current_x, current_y, current_z - distance_mm
        
        # 限位检查
        is_valid, error_msg = validate_arm_position(target_x, target_y, target_z)
        if not is_valid:
            return f"限位错误：{error_msg}。当前位置：({current_x:.1f}, {current_y:.1f}, {current_z:.1f})mm，无法执行该移动。"
        
        # 执行移动动作
        direction_map[direction](distance)
        
        return f"机械臂{direction}移动动作执行成功：距离{distance}厘米，目标位置({target_x:.1f}, {target_y:.1f}, {target_z:.1f})mm"
            
    except RuntimeError as e:
        return f"错误：{str(e)}"
    except Exception as e:
        logger.error(f"机械臂方向移动异常: {e}")
        return f"机械臂方向移动过程中出现错误：{str(e)}"

@FunctionRegistry.register_function('arm_gripper_control', '控制机械臂夹爪开合')
def arm_gripper_control_func(action: str) -> str:
    """
    控制机械臂夹爪开合
    
    Args:
        action: 动作类型 ("open"/"close")
        
    Returns:
        str: 执行结果描述
    """
    try:
        controller = get_arm_controller()
        
        action_map = {
            "open": controller.release_sth,
            "close": controller.catch_sth
        }
        
        if action not in action_map:
            return f"错误：无效的动作类型 '{action}'，支持的动作: {list(action_map.keys())}"
        
        # 执行夹爪控制动作
        action_map[action]()
        
        return f"机械臂夹爪控制动作执行成功：{action}"
            
    except RuntimeError as e:
        return f"错误：{str(e)}"
    except Exception as e:
        logger.error(f"机械臂夹爪控制异常: {e}")
        return f"机械臂夹爪控制过程中出现错误：{str(e)}"

@FunctionRegistry.register_function('arm_move_to_position', '控制机械臂移动到指定的绝对坐标位置(x,y,z)，适用于"移动到"、"到达位置"等绝对位置指令')
def arm_move_to_position_func(target_x: float, target_y: float, target_z: float) -> str:
    """
    控制机械臂移动到指定的绝对位置
    
    Args:
        target_x: 目标X坐标 (毫米)
        target_y: 目标Y坐标 (毫米)
        target_z: 目标Z坐标 (毫米)
        
    Returns:
        str: 执行结果描述
    """
    try:
        # 参数验证
        if not isinstance(target_x, (int, float)) or not isinstance(target_y, (int, float)) or not isinstance(target_z, (int, float)):
            return "错误：坐标参数必须是数字"
        
        # 严格限位检查 - 拒绝执行超限命令
        is_valid, error_msg = validate_arm_position(target_x, target_y, target_z)
        if not is_valid:
            return f"限位错误：{error_msg}。请调整目标位置到安全范围内。"
        
        controller = get_arm_controller()
        
        # 获取当前位置用于日志记录
        current_x, current_y, current_z = controller.get_current_position()
        
        # 执行绝对位置移动
        success = controller.arm_move(target_x, target_y, target_z)
        
        if success:
            return f"机械臂绝对移动成功：从位置({current_x:.1f}, {current_y:.1f}, {current_z:.1f})mm 移动到 ({target_x:.1f}, {target_y:.1f}, {target_z:.1f})mm"
        else:
            return f"机械臂绝对移动失败：目标位置({target_x:.1f}, {target_y:.1f}, {target_z:.1f})mm"
            
    except RuntimeError as e:
        return f"错误：{str(e)}"
    except Exception as e:
        logger.error(f"机械臂绝对移动异常: {e}")
        return f"机械臂绝对移动过程中出现错误：{str(e)}"

@FunctionRegistry.register_function('arm_status', '获取机械臂当前状态信息')
def arm_status_func() -> str:
    """
    获取机械臂当前状态信息
    
    Returns:
        str: 状态信息描述
    """
    try:
        controller = get_arm_controller()
        
        # 获取状态信息
        # 注意：RobotArm类没有专门的状态获取方法，这里返回一些基本信息
        result_parts = []
        result_parts.append("机械臂状态信息：")
        result_parts.append("- 控制器：已初始化")
        result_parts.append("- 串口：/dev/ttyS1")
        result_parts.append("- 波特率：115200")
        result_parts.append("- 当前位置：使用逆运动学坐标系")
        result_parts.append(f"- X轴限位：{ARM_LIMIT_X_MIN} ~ {ARM_LIMIT_X_MAX} mm")
        result_parts.append(f"- Y轴限位：{ARM_LIMIT_Y_MIN} ~ {ARM_LIMIT_Y_MAX} mm")
        result_parts.append(f"- Z轴限位：{ARM_LIMIT_Z_MIN} ~ {ARM_LIMIT_Z_MAX} mm")
        
        return "\n".join(result_parts)
        
    except RuntimeError as e:
        return f"错误：{str(e)}"
    except Exception as e:
        logger.error(f"获取机械臂状态异常: {e}")
        return f"获取机械臂状态过程中出现错误：{str(e)}"